﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using Microsoft.Win32.SafeHandles;

namespace NHook
{
	public class PageProtect : IDisposable
	{
		[DllImport("kernel32.dll", SetLastError = true)]
		static extern bool VirtualProtect(IntPtr lpAddress, uint dwSize,
		   uint flNewProtect, out uint lpflOldProtect);

		[DllImport("kernel32.dll", SetLastError = true)]
		static extern bool VirtualProtectEx(IntPtr processHandle, IntPtr lpAddress, uint dwSize,
		   uint flNewProtect, out uint lpflOldProtect);


		uint oldProtection;
		IntPtr ptr;
		uint _Size;
		SafeFileHandle _ProcessHandle;

		public PageProtect(IntPtr ptr, PageProtection protection, uint size)
			: this(null, ptr, protection, size)
		{
		}
		public PageProtect(SafeFileHandle processHandle, IntPtr ptr, PageProtection protection, uint size)
		{
			this.ptr = ptr;
			_Size = size;
			_ProcessHandle = processHandle;
			ThisVirtualProtect((uint)protection);
		}
		public PageProtect(IntPtr ptr, PageProtection protection)
			: this(ptr, protection, 0)
		{

		}


		private void ThisVirtualProtect(uint protection)
		{
			if(_ProcessHandle == null)
				Assert(VirtualProtect(ptr, _Size, (uint)protection, out oldProtection));
			else
				Assert(VirtualProtectEx(_ProcessHandle.DangerousGetHandle(), ptr, _Size, (uint)protection, out oldProtection));
		}

		
		private void Assert(bool result)
		{
			if(!result)
			{
				throw new Win32Exception();
			}
		}

		#region IDisposable Members

		public void Dispose()
		{
			ThisVirtualProtect(oldProtection);
		}

		#endregion
	}
	public enum PageProtection
	{
		ExecuteReadWrite = 0x40
	}
}
